/*  iImager (network disk imaging via iscsi)
	Copyright (C) 2009-2010 Yao Wei <njustyw@gmail.com>

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
*/

#include "Q_ISCSIFinder.h"
#include "trimps_forensic_global.h"
#include "iscdiscserv.h"

#define MAX_CONNECT_COUNT 6

Q_DECLARE_METATYPE(ISCSI_SERVICE_QEQUEST)

/** proxy pattern */
class ISCSIProxyThread : public QThread
{
public:

	enum SocketState{ Read, Write};

	ISCSIProxyThread(ISCSIFinder* parent, SocketState state): finder(parent), state(state){}

	virtual void run();

	void setWriteParams(const QVariant& data);

private:
	ISCSIFinder* finder;
	SocketState state;
	QVariant writeParam;
};

void ISCSIProxyThread::setWriteParams(const QVariant& data)
{
	writeParam = data;
}

void ISCSIProxyThread::run()
{
	/** Write to daemon socket */
	if(state == Write)
	{
		ISCSI_SERVICE_QEQUEST buffer = writeParam.value<ISCSI_SERVICE_QEQUEST>();
		finder->tcpSocket->write((const char*)&buffer, sizeof(buffer));
		finder->tcpSocket->flush();
		return;
	}

	/** Read from daemon socket */
	while (finder->tcpSocket->bytesAvailable() > 0) 
	{
		qint64 n = finder->tcpSocket->bytesAvailable();
		QByteArray datagram;
		ISCSI_SERVICE_RESPONCE* buffer = NULL;
		int sizeofData;
		int offset = 0;
		int type;

		DEBUG3_MSG(QString("We are receving %1 bytes data from socket").arg(n));
		
		datagram.resize(n);
		qint64 m = finder->tcpSocket->read(datagram.data(), n);
		if(m < 0)	return;

		while(1)
		{
			if(offset >= m)	break;

			buffer = (ISCSI_SERVICE_RESPONCE*)(datagram.data()+offset);
			sizeofData = buffer->sizeOfData;
			offset += sizeofData;
			type = buffer->type;

			if(type == DEVICE_STATUS)
			{
				ISCSI_DEVICE_STATUS* deviceStatus = NULL;
				deviceStatus = (ISCSI_DEVICE_STATUS*)buffer->data;
				int nDisk = (n - sizeof(ISCSI_SERVICE_RESPONCE)) / sizeof(ISCSI_DEVICE_STATUS);
				QStringList ipAddList, deviceAddList, targetAddList;
				QStringList ipRemoveList, deviceRemoveList, targetRemoveList;
				int addCount = 0, removeCount = 0;

				for(int i = 0; i < nDisk; i++)
				{
					QString ip = deviceStatus[i].ipaddr;
					QString target = deviceStatus[i].target;
					QString device = deviceStatus[i].device;

					if(deviceStatus[i].status == DEVICE_ADD)
					{
						if(!finder->liveAllDisks.contains(ip)) 
							finder->liveAllDisks[ip];

						if(!finder->liveAllDisks[ip].contains(device))
							finder->liveAllDisks[ip].insert(device, target);

						addCount ++;
						ipAddList << ip;
						deviceAddList << device;
						targetAddList << target;
					}
					else
					{
						if(finder->liveAllDisks.contains(ip)) 
						{
							if(finder->liveAllDisks[ip].contains(device))
								finder->liveAllDisks[ip].remove(device);
						}

						removeCount ++;
						ipRemoveList << ip;
						deviceRemoveList << device;
						targetRemoveList << target;
					}
					DEBUG3_MSG( ((deviceStatus[i].status == DEVICE_ADD)? QString("Add") : QString("Remove")) + "< " + ip + ", " + target + ", " + device + ">");
				}
				if(addCount)
					emit finder->diskAdd(ipAddList, deviceAddList, targetAddList);
				if(removeCount)
					emit finder->diskRemove(ipRemoveList, deviceRemoveList, targetRemoveList);
			}
			else if(type == HOST_STATUS)
			{
				ISCSI_RACK_STATUS* hostStatus = NULL;
				hostStatus = (ISCSI_RACK_STATUS*)buffer->data;
				
				DEBUG3_MSG( QString(hostStatus->ipaddr) + ((hostStatus->status == HOST_STATUS_CONNECTED)?" Connected":" Disconnected"));
				
				if(hostStatus->status == HOST_STATUS_CONNECTED)
				{
					emit finder->ipAddStatus(hostStatus->ipaddr, true);
				}
				else if(hostStatus->status == HOST_STATUS_DISCONNECTED)
				{
					emit finder->ipAddStatus(hostStatus->ipaddr, false);
				}
			}
		}
	}
}

ISCSIFinder::ISCSIFinder(QObject* parent /* = 0 */) : QObject(parent)
{
	DEBUG_MSG("Enter");

	socketConnectCount = 0;
	aboutToExit = false;

	tcpSocket = new QTcpSocket(this);
	connect(tcpSocket,SIGNAL(disconnected()),this,SLOT(serviceDisconnected()));
	connect(tcpSocket,SIGNAL(connected()),this,SLOT(serviceConnected()));
	connect(tcpSocket, SIGNAL(error(QAbstractSocket::SocketError)), this, SLOT(serviceError(QAbstractSocket::SocketError)));
	connect(tcpSocket, SIGNAL(readyRead()),this, SLOT(dataReceived()));

#ifdef Q_OS_WIN
	QProcess process;
	process.setReadChannelMode(QProcess::MergedChannels);
	process.start("sc query ImagerDaemon");
	process.waitForFinished(-1);
	QString outputStr = QString::fromLocal8Bit(process.readAllStandardOutput());
	if(outputStr.contains("RUNNING"));
	else if(outputStr.contains("STOPPED"))
		QProcess::execute("sc start ImagerDaemon");
#else
	QProcess::execute("killall daemon");
	//QProcess::execute(QCoreApplication::applicationDirPath()+"/daemon");
	QProcess* process = new QProcess;
	process->setReadChannelMode(QProcess::MergedChannels);
	process->start(QCoreApplication::applicationDirPath()+"/daemon");
	
#endif
}

ISCSIFinder::~ISCSIFinder()
{
	DEBUG_MSG("Leave");

	aboutToExit = true;
	closeDaemonConnection();
	tcpSocket->disconnectFromHost();
}

void ISCSIFinder::connectToService()
{
	tcpSocket->connectToHost("127.0.0.1", SERVICE_PORT);
}

void ISCSIFinder::closeDaemonConnection()
{
	DEBUG2_MSG("Disconnect Daemon Service");
	
	ISCSI_SERVICE_QEQUEST buffer;
	buffer.sizeOfData = sizeof(buffer);
	buffer.type = LOGOUT_ALL;

	tcpSocket->write((const char*)&buffer, sizeof(buffer));
	tcpSocket->flush();
	tcpSocket->waitForBytesWritten();
}

void ISCSIFinder::serviceError(QAbstractSocket::SocketError error)
{
	socketConnectCount ++;

	if(socketConnectCount >= MAX_CONNECT_COUNT)
		;//SpecialMessageBox::warnning(NULL, "", tr("Cannot connect to service, please start the service in advance."));
	else
		QTimer::singleShot(5000, this, SLOT(connectToService()));
}

void ISCSIFinder::serviceConnected()                     
{
	DEBUG2_MSG("Service is connected");
	
	socketConnectCount = 0;

	liveAllDisks.clear();
	emit serviceStatus(true);


	QStringList ipAddList, deviceAddList, targetAddList;
	ipAddList << QString("127.0.0.1") << QString("127.0.0.1");
#ifdef Q_OS_WIN
	deviceAddList << QString("E:\\sde.dd") << QString("D:\\en_VS2008_x86_x64.iso");
#else
	deviceAddList << QString("/mnt/e/sde.dd") << QString("/mnt/d/en_VS2008_x86_x64.iso");
#endif
	targetAddList << QString("disk3") << QString("disk8");
	emit diskAdd(ipAddList, deviceAddList, targetAddList);
}

void ISCSIFinder::serviceDisconnected()                     
{
	if(aboutToExit)	
		return;
	
	emit serviceStatus(false);

	QStringList ipRemoveList, deviceRemoveList, targetRemoveList;
	int removeCount = 0;
	for(QHash<QString,  QHash<QString, QString> >::iterator it = liveAllDisks.begin(); 
		it != liveAllDisks.end(); 
		it++)
	{
		for(QHash<QString, QString>::iterator it2 = it.value().begin(); 
			it2 != it.value().end();
			it2 ++)
		{
			removeCount ++;
			ipRemoveList << it.key();
			deviceRemoveList << it2.key();
			targetRemoveList << it2.value();
		}
		//it = liveAllDisks.erase(it);
	}

	if(removeCount)
		emit diskRemove(ipRemoveList, deviceRemoveList, targetRemoveList);
}

void ISCSIFinder::dataReceived()
{
	ISCSIProxyThread* proxyThread = new ISCSIProxyThread(this, ISCSIProxyThread::Read);
	connect(proxyThread, SIGNAL(finished()), proxyThread, SLOT(deleteLater()));
	proxyThread->start();
}

void ISCSIFinder::addIP(const QString &ip)
{
	DEBUG2_MSG("We are ready to connect " + ip);

	ISCSI_SERVICE_QEQUEST buffer;
	buffer.sizeOfData = sizeof(buffer);
	buffer.type = ADD_TARGET_PORTAL;
	strncpy(buffer.ipaddr, ip.toAscii().data(), sizeof(buffer.ipaddr));
	
	ISCSIProxyThread* proxyThread = new ISCSIProxyThread(this, ISCSIProxyThread::Write);
	proxyThread->setWriteParams(QVariant::fromValue(buffer));
	connect(proxyThread, SIGNAL(finished()), proxyThread, SLOT(deleteLater()));
	proxyThread->start();
}

void ISCSIFinder::removeIP(const QString& ip)
{
	DEBUG2_MSG("We will remove connection with " + ip);

	ISCSI_SERVICE_QEQUEST buffer;
	buffer.type = REMOVE_TARGET_PORTAL;
	buffer.sizeOfData = sizeof(buffer);
	strncpy(buffer.ipaddr, ip.toAscii().data(), sizeof(buffer.ipaddr));

	ISCSIProxyThread* proxyThread = new ISCSIProxyThread(this, ISCSIProxyThread::Write);
	proxyThread->setWriteParams(QVariant::fromValue(buffer));
	connect(proxyThread, SIGNAL(finished()), proxyThread, SLOT(deleteLater()));
	proxyThread->start();
}